import sys
import os, cv2
import numpy as np
import PyQt5.Qt as Qt
from PyQt5.QtGui import *
from PyQt5.QtCore import *
from PyQt5.QtWidgets import *
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "-1" # 禁用GPU
import keras
import tensorflow as tf

model = keras.models.load_model('mnist.h5')

def QPixmap2Mat(img : QPixmap):
    qimg = img.toImage()
    temp_shape = (qimg.height(), qimg.bytesPerLine() * 8 // qimg.depth())
    temp_shape += (4,)
    ptr = qimg.bits()
    ptr.setsize(qimg.byteCount())
    result = np.array(ptr, dtype=np.uint8).reshape(temp_shape)
    result = result[..., :3]
    return result

def expandImage(img):
    w, h = img.shape[0 : 2]
    size = max(w, h)
    x = (size - w) // 2
    y = (size - h) // 2
    image = np.zeros((size, size), dtype='uint8')
    for i in range(w):
        for j in range(h):
            image[i + x, j + y] = img[i, j]
    return image

def predictImage(img):
    img = cv2.resize(img, (28, 28))
    img = np.expand_dims(img, axis=0)
    return np.argmax(model.predict(img))

def predict(img: QPixmap):
    img = cv2.cvtColor(QPixmap2Mat(img), cv2.COLOR_RGB2GRAY)
    contours, _ = cv2.findContours(img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    string = ''
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        number = img[y : y + h, x : x + w]
        string += str(predictImage(expandImage(number)))
    return ' '.join(string)

class Canvas(QWidget):
    def __init__(self, parent):
        super(Canvas, self).__init__()
        self.setFixedSize(560, 560)
        self.setParent(parent)
        self.onDrawing = False
        self.setGeometry(10, 10, 560, 560)
        self.setStyleSheet("background-color:rgb(0,0,0);")
        self.pen = QPen()
        self.pen.setColor(QColor(255, 255, 255))
        self.pen.setWidth(30)
        self.pen.setJoinStyle(Qt.RoundJoin)
        self.pen.setCapStyle(Qt.RoundCap)
        self.buffer = QPixmap(560, 560)
        self.painter = QPainter()
        self.temp = QPainter()
        self.prev = QPoint()
        self.current = QPoint()

    def paintEvent(self, event):
        if self.onDrawing:
            self.temp.drawLine(self.prev, self.current)
        self.painter.begin(self)
        self.painter.setRenderHint(QPainter.Antialiasing)
        self.painter.setPen(self.pen)
        self.painter.drawPixmap(self.rect(), self.buffer)
        self.painter.end()
        return super().paintEvent(event)

    def mouseMoveEvent(self, event: QMouseEvent):
        if self.onDrawing:
            self.current = event.pos()
            self.repaint()
            self.prev = self.current
        return super().mouseMoveEvent(event)

    def mousePressEvent(self, event: QMouseEvent):
        if event.button() == Qt.LeftButton:
            self.temp.begin(self.buffer)
            self.temp.setRenderHint(QPainter.Antialiasing)
            self.temp.setPen(self.pen)
            self.onDrawing = True
            self.prev = self.current = event.pos()
            self.update()
        return super().mousePressEvent(event)

    def mouseReleaseEvent(self, event: QMouseEvent):
        if event.button() == Qt.LeftButton:
            self.update()
            self.temp.end()
            self.onDrawing = False
        return super().mouseReleaseEvent(event)

    def setWidth(self, w):
        self.pen.setWidth(w)

    def clear(self):
        self.buffer.fill(QColor(0, 0, 0))
        self.update()


class MainWindow(QMainWindow):
    def __init__(self):
        super(MainWindow, self).__init__()
        self.setWindowTitle("MNIST-GUI")
        self.setFixedSize(690, 580)
        self.canvas = Canvas(self)
        self.btn = QPushButton(self)
        self.btn.setGeometry(580, 10, 100, 50)
        self.btn.setText("Detect")
        self.clearBtn = QPushButton(self)
        self.clearBtn.setGeometry(580, 70, 100, 50)
        self.clearBtn.setText("Clear")
        self.label = QLabel(self)
        self.label.setGeometry(580, 130, 100, 50)
        self.label.setText("None")
        self.label.setStyleSheet("background-color: white;")
        self.label.setAlignment(Qt.AlignCenter)
        self.slider = QSlider(self)
        self.slider.setGeometry(580, 200, 100, 30)
        self.slider.setMinimum(1)
        self.slider.setMaximum(100)
        self.slider.setValue(30)
        self.slider.setOrientation(Qt.Horizontal)
        self.btn.clicked.connect(lambda: self.label.setText(predict(self.canvas.buffer)))
        self.clearBtn.clicked.connect(lambda: self.canvas.clear())
        self.slider.valueChanged.connect(
            lambda: self.canvas.setWidth(self.slider.value()))
        
if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    exit(app.exec())